import datasets.datasets_ws as datasets_ws
import cv2
import torch
import glob
import os
from os.path import join
import string
from torch.utils.data.dataset import Subset
import torch.utils.data as data

import numpy as np

class Dataset3DTiles(datasets_ws.BaseDataset):
    def __init__(self, context):
        config = context.config

        self.datasets_folder = config['data']['datasets_folder']
        dataset_name = config['data']['dataset_name']
        self.datasets_folder = self.datasets_folder  + "/" + dataset_name + "/"

        self._resize_flag = False
        if "resize" in context.config["data"]["preprocessing"].keys():
            self._resize_flag = True
            self.img_size = context.config["data"]["preprocessing"]["resize"]

        # self.database_paths = sorted(glob.glob(join(self.datasets_folder, "3DTiles/**", "*L22*.jpg"), recursive=True))
        self.database_paths = sorted(glob.glob(join(self.datasets_folder, "database/**", "*L22*.jpg"), recursive=True))

        self.database_num  = len(self.database_paths)

        self.images_paths = self.database_paths
        
        self._query_set = Dataset3DTilesQuery(context)

    def getDatabaseSubset(self):
        return Subset(self, list(range(self.database_num)))

    def getQuerySubset(self):
        return self._query_set
        # return Subset(self, list(range(self.database_paths)))

    def get_filename_byid(self,qid,fullpath = False):
        path = self.images_paths[qid]
        if fullpath:
            return path
        else:
            path = path.replace(self.datasets_folder,"")
            return path

    def __getitem__(self, index):
        imgpath = self.images_paths[index]
        input_image = self.read_image(imgpath)
        input_image = input_image.astype('float32')

        maskpath = imgpath.replace(".jpg", ".MASK.png")
        mask_image = self.read_image(maskpath)
        mask_image = mask_image.astype('float32')
        result = {
            "image":input_image,
            "mask":mask_image,
            "image_path": imgpath,
            "mask_path": maskpath,
        }
        return result, index
    
    def __len__(self):
        return super().__len__()
    

    def read_image(self, impath):
        """ Read image as grayscale and resize to img_size.
        Inputs
          impath: Path to input image.
          img_size: (W, H) tuple specifying resize size.
        Returns
          grayim: float32 numpy array sized H x W with values in range [0, 1].
        """
        grayim = cv2.imread(impath, 0)
        if grayim is None:
            raise Exception('Error reading image %s' % impath)
        # Image is resized via opencv.
       
        if self._resize_flag == True:
          interp = cv2.INTER_AREA
          img_size = self.img_size
          # size(width,height)
          grayim = cv2.resize(grayim, (img_size[0], img_size[1]), interpolation=interp)
        grayim = (grayim.astype('float32') / 255.)
        return grayim

class Dataset3DTilesQuery(data.Dataset):
    def __init__(self, context):
        config = context.config

        self.datasets_folder = config['data']['datasets_folder']
        dataset_name = config['data']['dataset_name']
        self.datasets_folder = self.datasets_folder  + "/" + dataset_name + "/"
 
        self.query_folder = config['data']['query_set']['query_dir']

        self._resize_flag = False
        self._to_float = True
        self._to_gray = False
        if "resize" in context.config["data"]["query_set"]["preprocessing"].keys():
            self._resize_flag = True
            self.img_size = context.config["data"]["query_set"]["preprocessing"]["resize"]

        if "to_float" in context.config["data"]["query_set"]["preprocessing"].keys():
            self._to_float = context.config["data"]["query_set"]["preprocessing"]["to_float"]
        if "to_gray" in context.config["data"]["query_set"]["preprocessing"].keys():
            self._to_gray = context.config["data"]["query_set"]["preprocessing"]["to_gray"]

        self.query_paths = sorted(glob.glob(join(self.datasets_folder, self.query_folder, "*.jpg"), recursive=True))
        self.query_num  = len(self.query_paths)
        
        # The format must be path/to/file/@utm_easting@utm_northing@...@.jpg
        self.query_coords = np.array([(path.split("@")[1], path.split("@")[2]) for path in self.query_paths]).astype(float)

        self.images_paths = self.query_paths
        
    def __getitem__(self, index):
        imgpath = self.images_paths[index]
        input_image = self.read_image_2(imgpath)
        
        result = {
            "image":input_image,
            "image_path":imgpath
        }
        return result,index
    
    def __len__(self):
        return self.query_num
    

    def get_filename_byid(self,qid,fullpath = False):
        path = self.images_paths[qid]
        if fullpath:
            return path
        else:
            path = path.replace(self.datasets_folder,"")
            return path

    def read_image_2(self, impath):
        """ Read image as grayscale and resize to img_size.
        Inputs
          impath: Path to input image.
          img_size: (W, H) tuple specifying resize size.
        Returns
          grayim: float32 numpy array sized H x W with values in range [0, 1].
        """
        img_data = cv2.imread(impath)
        if img_data is None:
            raise Exception('Error reading image %s' % impath)
    
        if self._to_gray:
            img_data = cv2.cvtColor(img_data,cv2.COLOR_BGR2GRAY)
        
        if self._to_float:
            img_data = (img_data.astype('float32') / 255.)

        # Image is resized via opencv.
        if self._resize_flag == True:
          interp = cv2.INTER_AREA
          img_size = self.img_size
          # size(width,height)
          img_data = cv2.resize(img_data, (img_size[0], img_size[1]), interpolation=interp)
        
        return img_data

    def read_image(self, impath):
        """ Read image as grayscale and resize to img_size.
        Inputs
          impath: Path to input image.
          img_size: (W, H) tuple specifying resize size.
        Returns
          grayim: float32 numpy array sized H x W with values in range [0, 1].
        """
        grayim = cv2.imread(impath, 0)
        if grayim is None:
            raise Exception('Error reading image %s' % impath)
        # Image is resized via opencv.
       
        if self._resize_flag == True:
          interp = cv2.INTER_AREA
          img_size = self.img_size
          # size(width,height)
          grayim = cv2.resize(grayim, (img_size[0], img_size[1]), interpolation=interp)
        grayim = (grayim.astype('float32') / 255.)
        return grayim